#include "HWJpegDecoder.h"
#include <string.h>
#include <stdio.h>
#include "FreeRTOS.h"
#include "cmsis_os.h"

void JPEG_GetDataCallback(JPEG_HandleTypeDef *hjpeg, uint32_t NbDecodedData);
void JPEG_DataReadyCallback(JPEG_HandleTypeDef *hjpeg, uint8_t *pDataOut, uint32_t OutDataLength);
void JPEG_ErrorCallback(JPEG_HandleTypeDef *hjpeg);
void JPEG_DecodeCpltCallback(JPEG_HandleTypeDef *hjpeg);

void HWJpegDecoderThread(void *arg);

HWJpegDecoder::HWJpegDecoder(JPEG_HandleTypeDef *hwJpegHandler)
{
	hwHandle = hwJpegHandler;

	// 使用MspDeInit保存对象
	HAL_JPEG_RegisterCallback(hwHandle, HAL_JPEG_MSPDEINIT_CB_ID, (void (*)(JPEG_HandleTypeDef *))this);

	// 注册回调函数
	HAL_JPEG_RegisterDataReadyCallback(hwHandle, JPEG_DataReadyCallback);
	HAL_JPEG_RegisterGetDataCallback(hwHandle, JPEG_GetDataCallback);
	HAL_JPEG_RegisterCallback(hwHandle, HAL_JPEG_ERROR_CB_ID, JPEG_ErrorCallback);
	HAL_JPEG_RegisterCallback(hwHandle, HAL_JPEG_DECODE_CPLT_CB_ID, JPEG_DecodeCpltCallback);

	taskENTER_CRITICAL();

	// 创建消息队列
	jpegQueue = xQueueCreate(10, sizeof(uint32_t));

	// 创建解码完成信号量
	decodeCpltSem = xSemaphoreCreateBinary();

	taskEXIT_CRITICAL();
}

HWJpegDecoder::~HWJpegDecoder()
{
	reset();	// 释放已分配的内存
	vQueueDelete(jpegQueue);		// 删除消息队列
	vSemaphoreDelete(decodeCpltSem);	// 删除信号量
	HAL_JPEG_UnRegisterCallback(hwHandle, HAL_JPEG_MSPDEINIT_CB_ID);	// 去除绑定的this

	// 注销绑定的回调函数
	HAL_JPEG_UnRegisterDataReadyCallback(hwHandle);
	HAL_JPEG_UnRegisterGetDataCallback(hwHandle);
	HAL_JPEG_UnRegisterCallback(hwHandle, HAL_JPEG_ERROR_CB_ID);
	HAL_JPEG_UnRegisterCallback(hwHandle, HAL_JPEG_DECODE_CPLT_CB_ID);
}

/**
 * @brief 获取JPEG图片信息
 */
HWJpegDecoder::ImageInfo &HWJpegDecoder::getImageInfo(void)
{
	return info;
}

/**
 * @brief 设置输入数据源(数组)
 */
void HWJpegDecoder::setInputDataSource(const uint8_t *data, uint32_t size)
{
	inputData = data;
	inputDataSize = size;
}

/**
 * @brief 设置输入数据源(文件路径)
 */
void HWJpegDecoder::setInputDataSource(const char *path)
{
	inputFilePath = path;
}

/**
 * @brief 获取解码的图片数据
 */
uint8_t *HWJpegDecoder::getImage(void)
{
	return outputBuffer;
}

/**
 * @brief 设置全尺寸图片数据输入
 * @note 若为true, 解码器会分配整个图片文件大小的内存空间并一次性读取图片数据。
 */
void HWJpegDecoder::setFullInput(bool en)
{
	fullInput = en;
}

/**
 * @brief 设置全尺寸图片数据输出
 * @note 若为true, 解码器会分配整个图片大小的内存空间，
 * @note 在解码完成后可以通过getImage()获取完整的图片数据。
 */
void HWJpegDecoder::setFullOutput(bool en)
{
	fullOutput = en;
}

/**
 * @brief 非全尺寸图片输出时，读取输出数据的回调函数
 * @note 全尺寸图片输出时，回调函数无效，请在解码后使用getImage()获取完整的图片数据。
 */
void HWJpegDecoder::setOutputCallback(OutputCb outputCb)
{
	outputCallback = outputCb;
}

/**
 * @brief 阻塞等待解码完成
*/
void HWJpegDecoder::waitDecodeFinish(void)
{
	xSemaphoreTake(decodeCpltSem, portMAX_DELAY);
}

/**
 * @brief 开启解码
 * @param onlyHeader 是否仅解码JPEG头部信息，包含图片宽高信息
 * @retval 若onlyHeader为true，则解析成功会返回NoData
 */
HWJpegDecoder::Error HWJpegDecoder::decode(bool onlyHeader)
{
	if (hwHandle == nullptr)
		return HWError;

	if (deviceBusyFlag)
		return DeviceBusy;

	Error ret = NoInputData;
	inputBufferSize = defaultInputBufferSize;	// 单次解码输入数据量

	do
	{
		if ((inputData == nullptr || inputDataSize < 32) && inputFilePath == nullptr)
			break;

		if (inputFilePath != nullptr) // 输入源有指定文件路径
		{
			inputFile = new FIL;
			memset(inputFile, 0, sizeof(FIL));

			if (f_open(inputFile, inputFilePath, FA_READ) != FR_OK)
			{
				ret = FileNoExist;
				break;
			}
			inputFrameSize = f_size(inputFile); // 获取文件大小

			if (fullInput)	
				inputBufferSize = inputFrameSize;		// 整个图像输入

			inputBuffer = new uint8_t[inputBufferSize];
			if (inputBuffer == nullptr)
			{
				ret = InputOutOfMemory;
				break;
			}

			size_t br;
			f_read(inputFile, (void *)inputBuffer, inputBufferSize, &br);

			if (fullInput) // 整个图片输入后续无需再使用文件
			{
				f_close(inputFile);
				delete inputFile;
				inputFile = nullptr;
			}
		}
		else		// 数组输入
		{
			inputFrameSize = inputDataSize; // 获取输入数据大小

			if (fullInput)	
				inputBufferSize = inputFrameSize;		// 整个图像输入

			inputBuffer = new uint8_t[inputBufferSize];
			if (inputBuffer == nullptr)
			{
				ret = InputOutOfMemory;
				break;
			}

			memcpy(inputBuffer, inputData, inputBufferSize);		// 拷贝数据
		}

		if (decodeHeader() != 0) // 解析JPEG头部
		{
			ret = InvalidImage;
			break;
		}

		if(onlyHeader)
		{
			ret = NoData;
			break;
		}

		// printf("info.w = %ld\r\n", info.width);
		// printf("info.h = %ld\r\n", info.height);
		// printf("info.s = %ld\r\n", info.chromaSubsampling);
		if (fullOutput)		// 全尺寸图片输出(分配整张图片所需的内存)
		{
			float bytePerPixel = 3; // YUV444
			if (info.chromaSubsampling == JPEG_422_SUBSAMPLING)
				bytePerPixel = 2;	// YUV422

			outputFrameSize = info.width * info.height * bytePerPixel;
			if (info.chromaSubsampling == JPEG_420_SUBSAMPLING)
				outputFrameSize /= 2; // YUV420
		}
		else		// 非全尺寸图片输出(使用defaultOutputBufferSize指定的内存空间进行输出缓冲)
		{
			info.MCUPerLine = info.width / info.MCUWidth;	// 1行MCU有多少个
			outputMCULine = defaultOutputBufferSize / (info.MCUPerLine * info.MCUSize);	// 计算缓冲区最多可存放几行
			if(outputMCULine == 0)		// 指定的缓冲区不足以放下一行MCU
			{
				ret = OutputOutOfMemory;
				break;
			}

			info.MCULines = info.height / info.MCUHeight;	// 需要多少行MCU
			if(info.height % info.MCUHeight != 0)
				info.MCULines += 1;

			if(info.MCULines < outputMCULine)	// 需要的行数比可存放的行数少
				outputMCULine = info.MCULines;

			outputFrameSize = outputMCULine * info.MCUPerLine * info.MCUSize;	// 单次输出需要使用的内存空间
		}

		uint32_t outputTransferSize = outputFrameSize;
		if(outputTransferSize > fullOutputSize)
			outputTransferSize = fullOutputSize;

		outputBuffer = new uint8_t[outputFrameSize];
		if (outputBuffer == nullptr)
		{
			ret = OutputOutOfMemory;
			break;
		}

		JPEGSourceAddress = (uint32_t)inputBuffer;		
		frameBufferAddress = (uint32_t)outputBuffer;
		inputFrameIndex = 0;
		outputFrameLines = 0;
		outputFrameLineIndex = 0;

		uint32_t inputTransferSize = fullInputSize;		// 单次输入传输数据量
		if(inputBufferSize < inputTransferSize)
			inputTransferSize = inputBufferSize;

		deviceBusyFlag = true;
		xQueueReset(jpegQueue); // 清空队列
		xQueueReset(decodeCpltSem);		// 清除信号量

		/* 启动JPEG解码 */
		if (HAL_JPEG_Decode_DMA(hwHandle, (uint8_t *)JPEGSourceAddress, inputTransferSize, (uint8_t *)frameBufferAddress, outputTransferSize) != HAL_OK)
		{
			ret = HWError;
			break;
		}

		taskENTER_CRITICAL();
		/* 启动解码线程 */
		xTaskCreate(HWJpegDecoderThread, "jpeg decode", 256, this, osPriorityNormal, nullptr);
		taskEXIT_CRITICAL();

		ret = NoError;
	} while (0);

	switch (ret)
	{
	case HWError:
			delete[] outputBuffer;
			outputBuffer = nullptr;
	case OutputOutOfMemory:
	case NoData:
	case InvalidImage:
			delete[] inputBuffer;
			inputBuffer = nullptr;
	case InputOutOfMemory:
		f_close(inputFile); // no break
	case FileNoExist:
		delete inputFile;
		inputFile = nullptr;
		break;
	default:;
	}

	return ret;
}

/**
 * @brief 复位解码器
 * @note 复位解码器将释放所有已分配的内存
 */
void HWJpegDecoder::reset(void)
{
	if (deviceBusyFlag != false)
		HAL_JPEG_Abort(hwHandle);

	if (inputFilePath != nullptr)	// 是文件源输入
	{
		if(inputFile != nullptr)	// 文件未关闭?
		{
			f_close(inputFile);
			delete inputFile;
			inputFile = nullptr;
		}
	}

	if(inputBuffer != nullptr)	
	{
		delete[] inputBuffer;
		inputBuffer = nullptr;
	}

	if(outputBuffer != nullptr)
	{
		delete[] outputBuffer;
		outputBuffer = nullptr;
	}

	inputFilePath = nullptr;
	inputFile = nullptr;
	inputDataSize = 0;
	inputData = nullptr;

	memset(&info, 0, sizeof(info));

	xQueueReset(jpegQueue);
	HWJpegDecoder::JpegMsg msg = HWJpegDecoder::JpegReset;
	xQueueSend(jpegQueue, &msg, portMAX_DELAY);		// 发送消息关闭解码线程(如果未关闭)

	deviceBusyFlag = false;
}

/**
 * @brief 图像解码线程
 */
void HWJpegDecoderThread(void *arg)
{
	HWJpegDecoder *decoder = (HWJpegDecoder *)arg;
	HWJpegDecoder::JpegMsg msg = HWJpegDecoder::DecodeInit;

	while (1)
	{
		xQueueReceive(decoder->jpegQueue, &msg, portMAX_DELAY);
		switch (msg)
		{
		case HWJpegDecoder::DecodeGetData:
			decoder->JPEG_GetDataProcess();
			break;
		case HWJpegDecoder::DecodeDataReady:
			decoder->JPEG_DataReadyProcess();
			break;
		case HWJpegDecoder::DecodeCplt:
			decoder->JPEG_DecodeCpltProcess(); // no break
		case HWJpegDecoder::JpegReset:
			vTaskDelete(NULL);
			break;
		case HWJpegDecoder::JpegError:
			decoder->JPEG_ErrorProcess();
			break;
		default:;
		}
	}
}

/**
 * @brief 获取输入数据回调函数
 */
void JPEG_GetDataCallback(JPEG_HandleTypeDef *hjpeg, uint32_t NbDecodedData)
{
	HWJpegDecoder *decoder = (HWJpegDecoder *)hjpeg->MspDeInitCallback;

	decoder->inputFrameIndex += NbDecodedData; // 更新已经解码的数据索引

	if (decoder->fullInput) // 全尺寸图片数据输入时，仅需要变更偏移地址
	{
		uint32_t inDataLength = 0;

		if (decoder->inputFrameIndex < decoder->inputFrameSize) // 若当前已经解码的数据小于总文件大小，继续解码
		{
			decoder->JPEGSourceAddress += NbDecodedData; // 更新解码数据偏移

			if ((decoder->inputFrameSize - decoder->inputFrameIndex) >= HWJpegDecoder::fullInputSize)
				inDataLength = HWJpegDecoder::fullInputSize; // 更新解码的数据大小(最大传输大小)
			else
				inDataLength = decoder->inputFrameSize - decoder->inputFrameIndex; // 更新解码的数据大小(剩余的大小)
		}

		HAL_JPEG_ConfigInputBuffer(hjpeg, (uint8_t *)decoder->JPEGSourceAddress, inDataLength); // 更新JPEG输入缓冲区
	}
	else
	{
		BaseType_t pxHigherPriorityTaskWoken = pdFALSE;
		HWJpegDecoder::JpegMsg msg = HWJpegDecoder::DecodeGetData; // 消息:读取解码数据块

		HAL_JPEG_Pause(decoder->hwHandle, JPEG_PAUSE_RESUME_INPUT);				 // 暂停JPEG数据输入
		xQueueSendFromISR(decoder->jpegQueue, &msg, &pxHigherPriorityTaskWoken); // 发送读数据消息
		portYIELD_FROM_ISR(pxHigherPriorityTaskWoken);
	}
}

/**
 * @brief 解码数据就绪回调函数
 */
void JPEG_DataReadyCallback(JPEG_HandleTypeDef *hjpeg, uint8_t *pDataOut, uint32_t OutDataLength)
{
	HWJpegDecoder *decoder = (HWJpegDecoder *)hjpeg->MspDeInitCallback;

	if (decoder->fullOutput) // 全尺寸图片数据输出时，仅需要变更偏移地址
	{
		decoder->frameBufferAddress += OutDataLength; // 更新图片输出数据偏移

		HAL_JPEG_ConfigOutputBuffer(hjpeg, (uint8_t *)decoder->frameBufferAddress, HWJpegDecoder::fullOutputSize); // 更新JPEG输出缓冲区
	}
	else
	{
		BaseType_t pxHigherPriorityTaskWoken = pdFALSE;
		HWJpegDecoder::JpegMsg msg = HWJpegDecoder::DecodeDataReady; // 消息:数据块解码完毕

		decoder->outputFrameLines = OutDataLength / decoder->info.MCUSize / decoder->info.MCUPerLine;	// 更新当前输出的行数

		HAL_JPEG_Pause(decoder->hwHandle, JPEG_PAUSE_RESUME_OUTPUT);			 // 暂停JPEG数据输出
		xQueueSendFromISR(decoder->jpegQueue, &msg, &pxHigherPriorityTaskWoken); // 发送消息
		portYIELD_FROM_ISR(pxHigherPriorityTaskWoken);
	}
}

/**
 * @brief 错误回调函数
 */
void JPEG_ErrorCallback(JPEG_HandleTypeDef *hjpeg)
{
	HWJpegDecoder *decoder = (HWJpegDecoder *)hjpeg->MspDeInitCallback;

	BaseType_t pxHigherPriorityTaskWoken = pdFALSE;
	HWJpegDecoder::JpegMsg msg = HWJpegDecoder::JpegError;

	xQueueSendFromISR(decoder->jpegQueue, &msg, &pxHigherPriorityTaskWoken);
	portYIELD_FROM_ISR(pxHigherPriorityTaskWoken);
}

/**
 * @brief 解码完成回调函数
 */
void JPEG_DecodeCpltCallback(JPEG_HandleTypeDef *hjpeg)
{
	HWJpegDecoder *decoder = (HWJpegDecoder *)hjpeg->MspDeInitCallback;

	BaseType_t pxHigherPriorityTaskWoken = pdFALSE;
	HWJpegDecoder::JpegMsg msg = HWJpegDecoder::DecodeCplt;

	xQueueSendFromISR(decoder->jpegQueue, &msg, &pxHigherPriorityTaskWoken);
	portYIELD_FROM_ISR(pxHigherPriorityTaskWoken);
}

void HWJpegDecoder::JPEG_GetDataProcess()
{
	size_t inDataLength = 0;

	if (inputFrameIndex < inputFrameSize)	// 已解码数据量小于总数据量
	{
		if(inputFilePath != nullptr)	// 文件输入
		{
			f_lseek(inputFile, inputFrameIndex);
			f_read(inputFile, (void *)inputBuffer, defaultInputBufferSize, &inDataLength);
		}
		else	// 数组输入
		{	
			inDataLength = inputFrameSize - inputFrameIndex;
			if(inDataLength > defaultInputBufferSize)
				inDataLength = defaultInputBufferSize;

			memcpy(inputBuffer, inputData, inDataLength);
		}

		if(inDataLength < 32UL)		// JPEG FIFO为32字节
		{
			memset((void *)(inputBuffer + inDataLength), 0, 32 - inDataLength);	// 剩余字节填充0
			inDataLength = 32;	// HAL_JPEG_Resume中没有对小于32字节的传输做处理，此处直接传输32字节
		}

		JPEGSourceAddress = (uint32_t)inputBuffer;
	}

	HAL_JPEG_ConfigInputBuffer(hwHandle, (uint8_t *)JPEGSourceAddress, inDataLength); // 更新JPEG输入缓冲区
	HAL_JPEG_Resume(hwHandle, JPEG_PAUSE_RESUME_INPUT);								  // 恢复数据输入
}

void HWJpegDecoder::JPEG_DataReadyProcess(void)
{
	if(outputCallback != nullptr)		// 调用回调函数
		outputCallback(info, outputBuffer, outputFrameLineIndex, outputFrameLines);

	outputFrameLineIndex += outputFrameLines;	// 标记下次输出的MCU行索引

	HAL_JPEG_ConfigOutputBuffer(hwHandle, (uint8_t *)frameBufferAddress, outputFrameSize); // 更新JPEG输出缓冲区
	HAL_JPEG_Resume(hwHandle, JPEG_PAUSE_RESUME_OUTPUT);	
}

void HWJpegDecoder::JPEG_DecodeCpltProcess(void)
{
	deviceBusyFlag = false;

	if (inputFilePath != nullptr)
	{
		if (inputFile != nullptr) // 文件未关闭(通常是非全尺寸输入的情况下)
		{
			f_close(inputFile);
			delete inputFile;
			inputFile = nullptr;
		}
	}

	delete[] inputBuffer;
	inputBuffer = nullptr;

	printf("decoder cplt\r\n");

	xSemaphoreGive(decodeCpltSem);	// 发送解码完成信号量
}

void HWJpegDecoder::JPEG_ErrorProcess(void)
{
	printf("jpeg error:%#lx\r\n", HAL_JPEG_GetError(hwHandle));
}

int HWJpegDecoder::decodeHeader(void)
{

#define MAKEUS(a, b) ((uint16_t)(((uint16_t)(a)) << 8 | ((uint16_t)(b))))

#define M_DATA 0x00
#define M_SOF0 0xc0
#define M_DHT 0xc4
#define M_SOI 0xd8
#define M_EOI 0xd9
#define M_SOS 0xda
#define M_DQT 0xdb
#define M_DNL 0xdc
#define M_DRI 0xdd
#define M_APP0 0xe0
#define M_APPF 0xef
#define M_COM 0xfe

	int ret = -5;
	int finished = 1;
	uint8_t id, ucHigh, ucLow;

	info.width = 0;
	info.height = 0;
	info.chromaSubsampling = -1;

	const uint8_t *pData = inputBuffer;

	while (finished)
	{
		id = *pData++;	// 段头0xff
		if (id != 0xff)
			break;

		id = *pData++;	// 段ID
		if (id >= M_APP0 && id <= M_APPF)
		{
			ucHigh = *pData++;
			ucLow = *pData++;
			pData += (long)(MAKEUS(ucHigh, ucLow) - 2);

			continue;
		}

		switch (id)
		{
		case M_SOI:
			break;

		case M_COM:
		case M_DQT:
		case M_DHT:
		case M_DNL:
		case M_DRI:
			ucHigh = *pData++;
			ucLow = *pData++;
			pData += (long)(MAKEUS(ucHigh, ucLow) - 2);
			break;

		case M_SOF0:	//基线式 - 图片宽高、颜色格式
			pData += 3;
			ucHigh = *pData++;
			ucLow = *pData++;
			info.height = MAKEUS(ucHigh, ucLow);
			ucHigh = *pData++;
			ucLow = *pData++;
			info.width = MAKEUS(ucHigh, ucLow);
			ret = 0;
			{
				uint8_t unitCnt = *pData++;
				uint8_t unitNum[3] = {0};
				for (uint8_t i = 0; i < unitCnt; i++)
				{
					uint8_t unitId = *pData++;
					if(unitId <= 3)
						unitNum[unitId - 1] = *pData++;

					pData += 1;
				}
				// 当前仅支持以下采样格式
				if (unitNum[0] == 0x22 && unitNum[1] == 0x22 && unitNum[2] == 0x22)		// CbCr全采样
				{
					info.chromaSubsampling = JPEG_444_SUBSAMPLING;	// YCbCr4:4:4
					info.MCUWidth = 8;
					info.MCUHeight = 8;
					info.MCUSize = 192;
				}
				else if (unitNum[0] == 0x22 && unitNum[1] == 0x12 && unitNum[2] == 0x12)	// CbCr垂直全采样，水平隔点采样
				{
					info.chromaSubsampling = JPEG_422_SUBSAMPLING;	// YCbCr4:2:2
					info.MCUWidth = 16;
					info.MCUHeight = 8;
					info.MCUSize = 256;
				}
				else if (unitNum[0] == 0x22 && unitNum[1] == 0x11 && unitNum[2] == 0x11)	// CbCr垂直水平隔点采样
				{
					info.chromaSubsampling = JPEG_420_SUBSAMPLING;	// YCbCr4:2:0
					info.MCUWidth = 16;
					info.MCUHeight = 16;
					info.MCUSize = 384;
				}
				else
					ret = -3;
			}

			finished = 0;
			break;

		case M_SOS:
		case M_EOI:
		case M_DATA:
			ret = -1;
			finished = 0;
			break;

		default:
			ucHigh = *pData++;
			ucLow = *pData++;
			pData += (long)(MAKEUS(ucHigh, ucLow) - 2);
			ret = -2;
			finished = 0;
			break;
		}
	}

	return ret;
}